from time import sleep

import yaml
from selenium import webdriver


class TestGetCookie:
    def setup(self):
        self.driver = webdriver.Chrome()

    def test_get_cookie(self):
        self.driver.get("https://work.weixin.qq.com/wework_admin/frame")
        sleep(15)
        cookies = self.driver.get_cookies()
        with open("./cookie.yaml","w") as f:
            yaml.safe_dump(cookies,f)

    def test_add_cookie(self):
        self.driver.get("https://work.weixin.qq.com/wework_admin/frame")
        cookie = yaml.safe_load(open("./cookie.yaml"))
        for c in cookie:
            self.driver.add_cookie(c)
        sleep(3)
        self.driver.get("https://work.weixin.qq.com/wework_admin/frame")